### ###########################################
### Estimates the ALR model
### Input is the formula selected from varsel
### Output is a model
### Priors unchanged because it's the same link
### ###########################################

here::i_am("newR/10_model_alr.R")

library(here)
library(brms)
library(tidyverse)
library(broom)
library(modelsummary)
library(knitr)

source("00_helpers.R")

dat <- readRDS("../working/tidy_model_data.rds")

load(here::here("working",
                "varsel_formula_alr.rdata"))


pvars <- c("Con_BE", "Lab_BE", "Lib_BE", "Nat_BE", "Oth2_BE")
for (p in pvars) { 
    dat[,p] <- replace(dat[,p],
                       dat[,p] <= 0,
                       1 / 40000)
}

### Make sure everything sums to 1
dat[,pvars] <- dat[,pvars] / rowSums(dat[,pvars])

dat$y <- with(dat,
              cbind(Con_BE, Lab_BE, Lib_BE, Nat_BE, Oth2_BE))

### Create a matrix as dep. var.
for (v in c("Lab_BE", "Lib_BE", "Nat_BE", "Oth2_BE")) {
    dat[,v] <- log(dat[,v] / dat$Con_BE)
}

varsel_f <- update(varsel_f,
                   mvbind(Lab_BE, Lib_BE, Nat_BE, Oth2_BE) ~ .)

### Establish priors
### Start with priors common to both models

common_priors <- c(prior("normal(0, 0.429)", class = "Intercept", resp = "LabBE"),
               prior("normal(-0.95, 0.584)", class = "Intercept", resp = "LibBE"),
               prior("normal(-2.457, 0.714)", class = "Intercept", resp = "NatBE"),
               prior("normal(-0.95, 0.584)", class = "Intercept", resp = "Oth2BE"))

### Labour Poll Change is common to both models
common_priors <- c(common_priors,
               prior("normal(0.241, 0.12)", coef = "LabPollChg_sc", resp = "LabBE"),
               prior("normal(0.06, 0.0905)", coef = "LabPollChg_sc", resp = "LibBE"),
               prior("normal(0.06, 0.0905)", coef = "LabPollChg_sc", resp = "NatBE"),
               prior("normal(0.06, 0.0905)", coef = "LabPollChg_sc", resp = "Oth2BE"))


### Now set up priors
### Start with the CV-minimizing model first
priors_min <- common_priors
coefs <- get_prior(varsel_f,
          data = dat) %>%
    filter(class == "b") %>%
    filter(coef != "") %>%
    pull(coef)

coefs <- setdiff(coefs,
                 "LabPollChg_sc")

resps <- get_prior(varsel_f,
          data = dat) %>%
    filter(class == "b") %>%
    filter(coef != "") %>%
    pull(resp) %>%
    unique()

for (k in coefs) {
    for (p in resps) {

        ### short_party <- sub("mu(...).?BE", "\\1", p)
        short_party <- sub("(...).?BE", "\\1", p)
        ## If it's a candidacy and the it concerns the party we're modelling
        if (grepl("Cand_BE", k) & grepl(short_party, k)) {
            prior_string <- "normal(0, 1)"
        } else {
            prior_string <- "normal(0, 0.25)"
        }
        
        priors_min <- c(priors_min,
                        set_prior(prior_string, class = "b",
                                  coef = k,
                                  resp = p))
    }
}


model_file <- here::here("working",
                         "alr_model.rds")

if (file.exists(model_file)) {
    mod <- readRDS(file = model_file)
} else {
    varsel_f <- update(varsel_f,
                       mvbind(Lab_BE, Lib_BE, Nat_BE, Oth2_BE) ~ .)

    ## Modelling the residual correlation is done by default.
    ## Which is nice
    mod <- brm(varsel_f,
               prior = priors_min,
               data = dat,
               control = list(adapt_delta = 0.95,
                              max_treedepth = 11),
               cores = 5, chains = 5)
    
    saveRDS(mod, file = model_file)
    
}

### ###################################################
### Start work on model evaluation
### ###################################################

### Get candidate matrix
candvars <- paste0(c("Con", "Lab", "Lib", "Nat", "Oth"),
                  "Cand_BE")
candmat <- as.matrix(dat[,candvars])

### Replace adjusted outcome values
y <- y0 <- dat$y
y[which(candmat == 0, arr.ind = TRUE)] <- 0

### Zero out fitted values
f0 <- fitted(mod, summary = FALSE)
f <- array(NA, dim = dim(f0) + c(0, 0, 1))
for (i in 1:dim(f0)[1]) {
    tmp <- f0[i,,]
    tmp <- cbind(1, exp(tmp))
    f[i,,] <- tmp / rowSums(tmp)
}

pp0 <- posterior_predict(mod)
pp <- array(NA, dim = dim(pp0) + c(0, 0, 1))
for (i in 1:dim(pp0)[1]) {
    tmp <- pp0[i,,]
    tmp <- cbind(1, exp(tmp))
    pp[i,,] <- tmp / rowSums(tmp)
}

f_old <- f
pp_old <- pp

## Zero out
for (i in 1:dim(f)[1]) {
    x <- f[i,,]
    x[which(candmat == 0, arr.ind = TRUE)] <- 0
    x <- x / rowSums(x)
    f[i,,] <- x
}

for (i in 1:dim(pp)[1]) {
    x <- pp[i,,]
    x[which(candmat == 0, arr.ind = TRUE)] <- 0
    x <- x / rowSums(x)
    pp[i,,] <- x
}

### Get different measures out
100 * get_mae3(f, y)
100 * get_mae3(f, y, type = "median")
100 * get_calibration3(pp, y)
100 * get_pcp3(f, y)
100 * get_brier3(f, y)

### What about predicted versus actual winner
party.names <- c("Con", "Lab", "Lib", "Nat", "Oth")

fbar <- apply(f, c(2, 3), mean)
pred_winner <- apply(fbar, 1, which.max)
pred_winner <- party.names[pred_winner]
actual_winner <- apply(dat[,pvars], 1, which.max)
actual_winner <- party.names[actual_winner]

